import datetime
import os
import sys
import threading
import zipfile
from pathlib import Path

import matlab.engine
from flask import Flask, jsonify, request, send_file
import sqlite3
import base64
import io
from flask_cors import CORS
from werkzeug.utils import secure_filename
from Spectrum_prediction import prediction
from NOMA import NOMA

app = Flask(__name__)
CORS(app)
UPLOAD_FOLDER = 'uploads'

os.makedirs(UPLOAD_FOLDER, exist_ok=True)
@app.route('/api/sate', methods=['POST'])
def receive_data():
    username = request.form.get('username')
    if not username:
        return jsonify({'error': 'Missing username'}), 400
    user_folder = os.path.join(UPLOAD_FOLDER, username)
    if not os.path.exists(user_folder):
        os.makedirs(user_folder)

    eng = matlab.engine.start_matlab()
    # 检查是否有文件上传
    if 'tleFile' not in request.files:
        return jsonify({'error': '未上传 TLE 文件'}), 400

    tle_file = request.files['tleFile']
    if tle_file.filename == '':
        return jsonify({'error': '未选择 TLE 文件'}), 400

    # 保存文件到本地
    file_path = os.path.join(user_folder, tle_file.filename)
    tle_file.save(file_path)

    satelliteIndex = request.form.get('satelliteIndex', type=int)
    start_time_str = request.form.get('startTime')  # "2024,1,7,14,5,0"
    end_time_str = request.form.get('endTime')  # "2024,1,7,20,5,0"
    MaxViewAngle =request.form.get('MaxViewAngle',type=int)
    # 获取地面站信息
    names = request.form.getlist('name[]')
    lats = request.form.getlist('lat[]')
    lons = request.form.getlist('lon[]')
    min_elevs = request.form.getlist('minElev[]')
    # 构建 groundStations 结构
    ground_stations = []
    for i in range(len(names)):
        station = {
            'latitude': float(lats[i]),
            'longitude': float(lons[i]),
            'minElevationAngle': float(min_elevs[i]),
            'name': names[i]
        }
        ground_stations.append(station)
    n = len(ground_stations)
    gs = eng.cell(n, 1)
    for i in range(n):
        station = ground_stations[i]
        gs[i] = eng.struct(
            'latitude', float(station['latitude']),
            'longitude', float(station['longitude']),
            'minElevationAngle', float(station['minElevationAngle']),
            'name', station['name']
        )
    # 打印接收到的数据（用于调试）
    r=eng.Satellite(file_path,gs,username,satelliteIndex, start_time_str, end_time_str,MaxViewAngle,nargout=4)
    eng.quit()
    access_csv_paths = r[0]
    aer_image_paths = r[1]
    doppler_image_paths = r[2]
    latency_image_paths = r[3]
    csv_files = [{'filename': os.path.basename(p), 'path': p} for p in access_csv_paths]

    # 处理图像文件：编码为 Base64
    def image_to_base64(path):
        with open(path, "rb") as img_file:
            return base64.b64encode(img_file.read()).decode("utf-8")

    aer_images = [{'filename': os.path.basename(p), 'base64': image_to_base64(p)} for p in aer_image_paths]
    doppler_images = [{'filename': os.path.basename(p), 'base64': image_to_base64(p)} for p in doppler_image_paths]
    latency_images = [{'filename': os.path.basename(p), 'base64': image_to_base64(p)} for p in latency_image_paths]

    # 构建响应数据
    response_data = {
        'csv_files': csv_files,
        'aer_images': aer_images,
        'doppler_images': doppler_images,
        'latency_images': latency_images
    }

    return jsonify(response_data)
    # 返回响应


@app.route('/api/download')
def download_file():
    file_path = request.args.get('file')
    if not file_path or not os.path.exists(file_path):
        return jsonify({"error": "File not found"}), 404

    return send_file(file_path, as_attachment=True)

@app.route('/api/upload/spectrum', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400

    file = request.files['file']
    username = request.form.get('username')

    if not username:
        return jsonify({'error': 'Missing username'}), 400

    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400

    # 构建用户专属目录
    user_dir = os.path.join('uploads', username)
    if not os.path.exists(user_dir):
        os.makedirs(user_dir)

    # 保存文件到指定路径
    file_path = os.path.join(user_dir, file.filename)
    file.save(file_path)

    return jsonify({'filepath': file_path}), 200


@app.route('/api/predict', methods=['POST'])
def predict():
    # 自动识别 content-type 类型
    if request.content_type == 'application/json':
        data = request.get_json()
    else:
        data = request.form
    username = data.get('username')
    print(username)
    if not username:
        return jsonify({'error': 'Missing username'}), 400

    # 使用 data 来获取参数
    TimeStep = int(data.get('TimeStep', 50))
    PredStep = int(data.get('PredStep', 1))
    Epochs = int(data.get('Epochs', 100))
    LearningRate = float(data.get('LearningRate', 0.001))
    StartFreq = int(data.get('StartFreq', 18))
    FreqNum = int(data.get('FreqNum', 15))
    BatchSize = int(data.get('BatchSize', 64))
    filename = data.get('filename')

    if not filename:
        return jsonify({'error': 'Missing filename'}), 400

    file_path = filename
    print(filename)
    if not os.path.exists(file_path):
        return jsonify({'error': 'File not found'}), 404
    user_spectrum_dir = os.path.join(username, 'spectrum')
    os.makedirs(user_spectrum_dir, exist_ok=True)


    # 调用预测函数
    result = prediction(
        FreqNum=FreqNum,
        StartFreq=StartFreq,
        TimeStep=TimeStep,
        PredStep=PredStep,
        BatchSize=BatchSize,
        Epochs=Epochs,
        LearningRate=LearningRate,
        Patience=7,
        filepath=file_path,
        output_dir=user_spectrum_dir
    )
    image_path=result["image_path"]
    images = []
    with open(image_path, 'rb') as img_file:
            img_str = base64.b64encode(img_file.read()).decode('utf-8')
            # 将图片信息添加到列表中
            images.append({
                'filename': image_path,
                'image_base64': img_str
            })
    return jsonify({
        'success': True,
        'result': result,
        'images': images
    }), 200






@app.route('/api/noma', methods=['POST'])
def run_matlab():
    data = request.json
    print("Received data:", data)
    hap_count = data.get('hap_count')
    p_hap = data.get('p_hap')
    username = data.get('username')  # 新增字段

    if not all([hap_count, p_hap, username]):
        return jsonify({"error": "Missing required parameters"}), 400

    # 构造专属输出路径
    output_dir = os.path.join(username, 'nomaoutputImages')

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    image_paths=NOMA(M=hap_count, P_hap=p_hap, output_dir=output_dir)

    images = []
    for img_path in image_paths:
        with open(img_path, 'rb') as img_file:
            img_str = base64.b64encode(img_file.read()).decode('utf-8')
        images.append({
            'filename': os.path.basename(img_path),
            'image_base64': img_str
        })
    return jsonify(images), 200
@app.route('/api/register', methods=['POST'])
def register():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')

    if not username or not password:
        return jsonify({"error": "Missing username or password"}), 400

    try:
        conn = sqlite3.connect('example.db')
        c = conn.cursor()
        c.execute('''
                    CREATE TABLE IF NOT EXISTS users (
                        id INTEGER PRIMARY KEY AUTOINCREMENT,
                        username TEXT UNIQUE NOT NULL,
                        password TEXT NOT NULL
                    )
                ''')
        c.execute("INSERT INTO users (username, password) VALUES (?, ?)", (username, password))
        conn.commit()
        conn.close()
        return jsonify({"message": "注册成功", "username": username}), 201
    except sqlite3.IntegrityError:
        return jsonify({"error": "用户名已存在"}), 400
    except Exception as e:
        return jsonify({"error": str(e)}), 500
@app.route("/api/login", methods=['POST'])
def login():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')
    if not username or not password:
        return jsonify({"error": "用户名或密码不能为空"}), 400

    try:
        conn = sqlite3.connect('example.db')
        c = conn.cursor()

        # 查询用户是否存在
        c.execute("SELECT * FROM users WHERE username=?", (username,))
        user = c.fetchone()

        if not user:
            return jsonify({"error": "用户不存在"}), 404

        db_id, db_username, db_password = user

        if password == db_password:
            return jsonify({
                "message": "登录成功",
                "username": db_username
            }), 200
        else:
            return jsonify({"error": "密码错误"}), 401

    except Exception as e:
        return jsonify({"error": str(e)}), 500
    finally:
        conn.close()

CLASS_NAMES = {
    0: "BPSK",
    1: "QPSK",
    2: "8PSK",
    3: "OQPSK",
    4: "16QAM",
    5: "2ASK",
    6: "16APSK",
    7: "2FSK",
    8: "FM",
    9: "DSB-AM",
    10: "Single",
    11: "Chirp",
    12: "Comb"
}
NC = len(CLASS_NAMES)

def find_image_dirs(root_path):
    """
    尝试自动查找 train 和 valid 图像路径
    """
    possible_train = [
        os.path.join(root_path, 'train', 'images', 'images'),
        os.path.join(root_path, 'train', 'images'),
        os.path.join(root_path, 'images', 'train'),
    ]
    possible_valid = [
        os.path.join(root_path, 'valid', 'images', 'images'),
        os.path.join(root_path, 'valid', 'images'),
        os.path.join(root_path, 'images', 'val'),
        os.path.join(root_path, 'images', 'valid'),
    ]

    train_path = None
    for path in possible_train:
        if os.path.isdir(path):
            train_path = path
            break

    val_path = None
    for path in possible_valid:
        if os.path.isdir(path):
            val_path = path
            break

    return {
        'train': train_path,
        'valid': val_path,
        'train_exists': train_path is not None,
        'val_exists': val_path is not None
    }
@app.route('/api/dcgz/upload', methods=['POST'])
def upload_dataset():
    print("Received form data:", request.form)
    print("Received files:", request.files)

    if 'file' not in request.files:
        return jsonify({'error': '没有文件'}), 400

    file = request.files['file']

    if file.filename == '':
        return jsonify({'error': '未选择文件'}), 400

    username = request.form.get('username', '')
    upload_type = request.form.get('type', 'train')

    if not username:
        return jsonify({'error': '缺少用户名参数'}), 400

    if '.' not in file.filename or file.filename.rsplit('.', 1)[1].lower() != 'zip':
        return jsonify({'error': '仅支持 .zip 格式'}), 400

    # 构建路径
    user_upload_dir = os.path.join(UPLOAD_FOLDER, username)
    os.makedirs(user_upload_dir, exist_ok=True)

    filename = secure_filename(file.filename)
    file_path = os.path.join(user_upload_dir, filename)

    try:
        # 删除旧文件
        if os.path.exists(file_path):
            os.remove(file_path)

        # 保存上传的 zip 文件
        file.save(file_path)

        # 解压路径
        dataset_name = os.path.splitext(filename)[0]
        extract_base_dir = os.path.join(user_upload_dir, 'dcgz', 'datasets', upload_type)
        extract_path = os.path.join(extract_base_dir, dataset_name)

        os.makedirs(extract_base_dir, exist_ok=True)

        # 解压 ZIP 文件
        with zipfile.ZipFile(file_path, 'r') as zip_ref:
            zip_ref.extractall(extract_path)

        # 查找真实数据集根目录（可能存在嵌套）
        real_dataset_root = os.path.join(extract_path, dataset_name)
        if not os.path.exists(real_dataset_root):  # 防止多层嵌套
            real_dataset_root = extract_path

        # 自动查找训练/验证图像路径
        base_path = "serve\\"

        if upload_type == 'train':
            paths = find_image_dirs(real_dataset_root)
            print(paths)

            if not paths['train_exists']:
                return jsonify({'error': f"未找到训练图像目录: {paths['train']}"}), 400
            if not paths['val_exists']:
                return jsonify({'error': f"未找到验证图像目录: {paths['valid']}"}), 400

            # 构建相对路径用于 YAML 中（相对于 path）
            train_rel = os.path.relpath(paths['train'], real_dataset_root).replace('\\', '/')
            val_rel = os.path.relpath(paths['valid'], real_dataset_root).replace('\\', '/')

            real_dataset_root = os.path.join(base_path, real_dataset_root)

            # 生成 YAML 内容
            yaml_content = f"""# YOLOv5 dataset for {dataset_name}
path: {real_dataset_root}  # 数据集根目录
train: {train_rel}  # 训练图像路径
val: {val_rel}      # 验证图像路径

# 类别数量
nc: {NC}

# 类别名称
names:
"""
            for k, v in CLASS_NAMES.items():
                yaml_content += f"  {k}: {v}\n"

            # 生成 YAML 文件路径
            yaml_file_path = os.path.join(extract_base_dir, f'{dataset_name}.yaml')
            with open(yaml_file_path, 'w', encoding='utf-8') as f:
                f.write(yaml_content)

            return jsonify({
                'message': '训练数据集上传并解压成功',
                'type': upload_type,
                'dataset_name': dataset_name,
                'real_dataset_root': real_dataset_root,
                'yaml_file': yaml_file_path
            })

        elif upload_type == 'predict':
            return jsonify({
                'message': '预测数据集上传并解压成功',
                'type': upload_type,
                'dataset_name': dataset_name,
                'extracted_path': real_dataset_root.replace('\\', '/'),
            })

        else:
            return jsonify({'error': '不支持的上传类型'}), 400

    except Exception as e:
        return jsonify({'error': f'文件操作失败: {str(e)}'}), 500



FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
def get_model_and_metrics_paths(project_path):
    """
    根据 project_path 查找训练结果文件
    :param project_path: Path 对象，例如 runs/train/exp
    :return: 包含模型和图像路径的字典
    """
    model_dir = project_path / "weights"
    result = {}

    # best_model 路径
    best_model_path = model_dir / "best.pt"
    result['best_model'] = str(best_model_path) if best_model_path.exists() else None

    # 图像文件路径
    p_curve_path = project_path / "P_curve.png"
    r_curve_path = project_path / "R_curve.png"
    results_png_path = project_path / "results.png"

    # 工具函数：读取图像并转为 base64
    def read_image_b64(path):
        if path.exists():
            with open(path, "rb") as f:
                return base64.b64encode(f.read()).decode("utf-8")
        return None

    # 三张图的 base64 数据
    result['p_curve'] = read_image_b64(p_curve_path)
    result['r_curve'] = read_image_b64(r_curve_path)
    result['results_png'] = read_image_b64(results_png_path)
    return result

def run_yolo_training(yaml_path, epochs=50, batch_size=8, username=None):
    from train import create_opt, main  # 假设你封装了 create_opt 函数

    if not username:
        raise ValueError("Username is required for saving training results.")

    # 构造项目路径：/path/to/your/training_root/username/dcgz/train_out
    project = Path(username) / "dcgz" / "train_out"
    name = "exp"  # 实验名称，可以自定义或从前端传入

    # 创建目录（如果不存在）
    project.mkdir(parents=True, exist_ok=True)
    full_project_path = project / name
    opt = create_opt(
        data=Path(yaml_path),
        epochs=int(epochs),
        batch_size=int(batch_size),
        project=project,
        name=name
    )

    main(opt)
    result = get_model_and_metrics_paths(full_project_path)
    training_results[username] = result
    return result
training_results = {}
@app.route('/api/dcgz/train', methods=['POST'])
def dcgz_start_training():
    data = request.json
    print("Received training data:", data)
    yaml_path = data.get('yaml_path')
    epochs = data.get('epochs', 50)
    batch_size = data.get('batch_size', 8)
    username = data.get('username')
    if not username:
        return jsonify({'status': 'error', 'message': '缺少用户名信息'}), 400

    if not yaml_path or not os.path.exists(yaml_path):
        return jsonify({'status': 'error', 'message': 'YAML 文件路径不存在'}), 400
    # 启动后台线程进行训练
    thread = threading.Thread(
        target=run_yolo_training,
        kwargs={
            'yaml_path': yaml_path,
            'epochs': epochs,
            'batch_size': batch_size,
            'username': username
        }
    )
    thread.start()
    output_dir = Path(username) / "dcgz" / "train_out"
    return jsonify({
        'status': 'success',
        'message': '训练已在后台启动',
        'output_dir': str(output_dir)
    })
@app.route('/api/dcgz/train/result', methods=['GET'])
def get_training_result():
    username = request.args.get('username')
    if not username:
        return jsonify({'error': '缺少用户名参数'}), 400

    result = training_results.get(username)
    if not result:
        return jsonify({'error': '暂无可用结果，请等待训练完成'}), 404

    return jsonify({
        'status': 'success',
        'result': result
    })

if __name__ == '__main__':
    app.run(debug=True)
